import numpy as np
from scipy.fftpack import fft
import random

from features.SpeechFeatureMeta import SpeechFeatureMeta

class SpecAugment(SpeechFeatureMeta):
    '''
    复现谷歌SpecAugment数据增强特征算法，基于Spectrogram语谱图基础特征
    '''
    def __init__(self,framerate=16000,timewindow=25,timeshift=10):
        self.time_window=timewindow
        self.window_length=int(framerate/1000*self.time_window)
        self.timeshift=timeshift
        self.x=np.linspace(0,self.window_length-1,self.window_length,dtype=np.int16)
        self.w=0.54-0.46*np.cos(2*np.pi*self.x/(self.window_length-1))
        super().__init__(framerate)

    def run(self,wavsignal,samplerate=16000):
        self.framerate=samplerate
        range0_end=int(len(wavsignal[0])/self.framerate*1000-self.time_window)//self.timeshift+1
        data_input=np.zeros((range0_end,self.window_length//2),dtype=np.float)
        data_line=np.zeros((1,self.window_length),dtype=np.float)
        for i in range(0,range0_end):
            p_start=i*int(self.framerate/1000*self.timeshift)
            p_end=p_start+self.window_length
            data_line=wavsignal[0,p_start:p_end]
            data_line=data_line*self.w
            data_line=np.abs(fft(data_line))
            data_input[i]=data_line[0:self.window_length//2]
        data_input=np.log(data_input+1)
        mode=random.randint(1,100)
        h_start=random.randint(1,data_input.shape[0])
        h_width=random.randint(1,100)
        v_start=random.randint(1,data_input.shape[1])
        v_width=random.randint(1,100)
        if mode<=60:
            pass
        elif 60<mode<=75:
            data_input[h_start:h_start+h_width,:]=0
        elif 75<mode<=90:
            data_input[:,v_start:v_start+v_width]=0
        else:
            data_input[h_start:h_start+h_width,v_start:v_start+v_width]=0
        return data_input
